import torch
import torch.nn as nn

class AngleSolver(nn.Module):
    
    def __init__(self,):
        super(AngleSolver, self).__init__()
        
        self.layer1 = nn.Sequential(
            nn.Linear(16, 32),
            nn.Linear(32, 64),
            nn.LayerNorm(64),
            nn.LeakyReLU(),
            nn.Dropout(0.1),

            nn.Linear(64, 128),
            nn.LayerNorm(128),
            nn.LeakyReLU(),
            nn.Dropout(0.1),

            nn.Linear(128, 64),
            nn.LayerNorm(64),
            nn.LeakyReLU(),

            nn.Linear(64, 32),
            nn.LayerNorm(32),
            nn.LeakyReLU(),

            nn.Linear(32, 16),
            nn.LayerNorm(16),
            nn.LeakyReLU(),

            nn.Linear(16, 6),
        )
    
    def forward(self, x):
        
        return self.layer1(x)
    